
import numpy as np
import torch
from PIL import Image
from medpy import metric
from scipy.ndimage import zoom
import torch.nn as nn
import SimpleITK as sitk
import torchvision.transforms.functional as TF
import random
from tqdm import tqdm
from matplotlib import pyplot as plt




class DiceLoss(nn.Module):
    def __init__(self, n_classes):
        super(DiceLoss, self).__init__()
        self.n_classes = n_classes

    def _one_hot_encoder(self, input_tensor):
        tensor_list = []
        for i in range(self.n_classes):
            temp_prob = input_tensor == i  # * torch.ones_like(input_tensor)
            tensor_list.append(temp_prob.unsqueeze(1))
        output_tensor = torch.cat(tensor_list, dim=1)
        return output_tensor.float()

    def _dice_loss(self, score, target):
        target = target.float()
        smooth = 1e-5
        intersect = torch.sum(score * target)
        y_sum = torch.sum(target * target)
        z_sum = torch.sum(score * score)
        loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
        loss = 1 - loss
        return loss

    def forward(self, inputs, target, weight=None, softmax=False):
        if softmax:
            inputs = torch.softmax(inputs, dim=1)
        target = self._one_hot_encoder(target)
        if weight is None:
            weight = [1] * self.n_classes
        assert inputs.size() == target.size(), 'predict {} & target {} shape do not match'.format(inputs.size(), target.size())
        class_wise_dice = []
        loss = 0.0
        # print("inputs:{}".format(inputs))
        # print("target:{}".format(target))
        # X[n0,n1] 表示第 0 维 取第n0 个元素 ，第 1 维取第 n1 个元素
        # dice = self._dice_loss(inputs, target)
        # print("dice:{}".format(dice))
        # class_wise_dice.append(1.0 - dice.item())
        # loss += dice * weight[1]
        for i in range(0, self.n_classes):
            # print("inputs[:, i]:{}".format(inputs[:, i]))
            # print("target[:, i]:{}".format(target[:, i]))
            dice = self._dice_loss(inputs[:, i], target[:, i])
            # print("dice:{}".format(dice))
            '''
                报错
                block: [23,0,0], thread: [489,0,0] Assertion `t >= 0 && t < n_classes` failed.

            '''
            class_wise_dice.append(1.0 - dice.item())
            loss += dice * weight[i]
        return loss / self.n_classes

"""
    换损失函数  -------------没用上-------------------
"""
class BCELoss(nn.Module):
    def __init__(self):
        super(BCELoss, self).__init__()
        self.bceloss = nn.BCELoss()

    def forward(self, pred, target):
        size = pred.size(0)  # (D, C, H, W)
        pred_ = pred.view(size, -1)  # (D, 3*H*W) (24,150528)
        # target = target.expand(24, 3, 224, 224)
        target_ = target.view(size, -1)  # (D, 1*H*W) (24,150528)


        # z_sum = torch.sum(pred * pred)
        # y_sum = torch.sum(target * target)

        return self.bceloss(pred_, target_)


# class DiceLoss(nn.Module):
#     def __init__(self):
#         super(DiceLoss, self).__init__()
#
#     def forward(self, pred, target):
#         smooth = 1
#         size = pred.size(0)
#
#         pred_ = pred.view(size, -1)
#         target_ = target.view(size, -1)
#         intersection = pred_ * target_
#         dice_score = (2 * intersection.sum(1) + smooth)/(pred_.sum(1) + target_.sum(1) + smooth)
#         dice_loss = 1 - dice_score.sum()/size
#
#         return dice_loss


class BceDiceLoss(nn.Module):
    def __init__(self, wb=1, wd=1):
        super(BceDiceLoss, self).__init__()
        self.bce = BCELoss()
        self.dice = DiceLoss()
        self.wb = wb
        self.wd = wd

    def forward(self, pred, target):
        bceloss = self.bce(pred, target)
        diceloss = self.dice(pred, target)

        loss = self.wd * diceloss + self.wb * bceloss
        return loss


def calculate_metric_percase(pred, gt):
    pred[pred > 0] = 1
    gt[gt > 0] = 1
    if pred.sum() > 0 and gt.sum()>0:
        dice = metric.binary.dc(pred, gt)
        hd95 = metric.binary.hd95(pred, gt)
        return dice, hd95
    elif pred.sum() > 0 and gt.sum()==0:
        return 1, 0
    else:
        return 0, 0


def test_single_volume(image, label, net, classes, patch_size=[256, 256], test_save_path=None, case=None, z_spacing=1):
    # 将image和label进行处理，将其从张量转换为NumPy数组，并移除维度为1的维度。
    image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    if len(image.shape) == 3:  # 维度为3，表示为切片数据
        prediction = np.zeros_like(label)  # 与label相同的全零数组用来存储预测结果
        for ind in range(image.shape[0]):
            slice = image[ind, :, :]  # 选择当前切片
            x, y = slice.shape[0], slice.shape[1]  # 获取当前切片的高度和宽度
            if x != patch_size[0] or y != patch_size[1]:  # 如果当前切片的尺寸与patch_size不匹配，则调整切片的尺寸使其与patch_size相匹配。
                slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=3)  # previous using 0
                #  转换为张量，并对其添加额外的维度，并将其发送到GPU上。
            input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()  # （1,1,224,224） 转换为张量 添加额外的维度
            input = input.repeat(1, 4, 1, 1)
            # input = input.numpy().astype(np.float32)

            net.eval()
            with torch.no_grad():
                outputs = net(input)  # （1,9,224,224）
                out = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0)  # （224,224）
                out = out.cpu().detach().numpy()
                # label_val = label[ind]


                if x != patch_size[0] or y != patch_size[1]:
                    pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
                else:
                    pred = out
                prediction[ind] = pred
    else:  # 维度不为3，则直接对整个图像进行预测
        input = torch.from_numpy(image).unsqueeze(0).unsqueeze(0).float().cuda()
        net.eval()
        with torch.no_grad():
            out = torch.argmax(torch.softmax(net(input), dim=1), dim=1).squeeze(0)
            prediction = out.cpu().detach().numpy()
    metric_list = []
    for i in range(1, classes):  # 计算预测结果和标签之间的评估指标
        metrics = calculate_metric_percase(prediction == i, label == i)
        metric_list.append(metrics)


    if test_save_path is not None:
        img_itk = sitk.GetImageFromArray(image.astype(np.float32))
        prd_itk = sitk.GetImageFromArray(prediction.astype(np.float32))
        lab_itk = sitk.GetImageFromArray(label.astype(np.float32))
        img_itk.SetSpacing((1, 1, z_spacing))
        prd_itk.SetSpacing((1, 1, z_spacing))
        lab_itk.SetSpacing((1, 1, z_spacing))
        sitk.WriteImage(prd_itk, test_save_path + '/'+case + "_pred.nii.gz")
        sitk.WriteImage(img_itk, test_save_path + '/'+ case + "_img.nii.gz")
        sitk.WriteImage(lab_itk, test_save_path + '/'+ case + "_gt.nii.gz")

    return metric_list

"""
    修改
"""
# def test_single_volume(image, label, net, classes, patch_size=[256, 256], test_save_path=None, case=None,
#                        z_spacing=1):
#     image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    # _, x, y = image.shape
    # if x != patch_size[0] or y != patch_size[1]:
    #     # 缩放图像符合网络输入
    #     image = zoom(image, (1, patch_size[0] / x, patch_size[1] / y), order=3)
    # input = torch.from_numpy(image).unsqueeze(0).float().cuda()
    # net.eval()
    # with torch.no_grad():
    #     out = torch.argmax(torch.softmax(net(input), dim=1), dim=1).squeeze(0)
    #     out = out.cpu().detach().numpy()
    #     if x != patch_size[0] or y != patch_size[1]:
    #         # 缩放图像至原始大小
    #         prediction = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
    #     else:
    #         prediction = out
    #
    # metric_list = []
    # for i in range(1, classes):
    #     metric_list.append(calculate_metric_percase(prediction == i, label == i))

    # 输出方便观看的图片 a1,a2,a3分别代表RGB三个通道 a1[a1==2]=0代表R通道中输出结果为2的赋值0，a2[a2==2]=255代表G通道中输出结果为2的赋值255
    # if test_save_path is not None:
    #     a1 = copy.deepcopy(prediction)
    #     a2 = copy.deepcopy(prediction)
    #     a3 = copy.deepcopy(prediction)
    #
    #     a1[a1 == 1] = 255
    #     a1[a1 == 2] = 0
    #     a1[a1 == 3] = 255
    #     a1[a1 == 4] = 20
    #
    #     a2[a2 == 1] = 255
    #     a2[a2 == 2] = 255
    #     a2[a2 == 3] = 0
    #     a2[a2 == 4] = 10
    #
    #     a3[a3 == 1] = 255
    #     a3[a3 == 2] = 77
    #     a3[a3 == 3] = 0
    #     a3[a3 == 4] = 120
    #
    #     a1 = Image.fromarray(np.uint8(a1)).convert('L')
    #     a2 = Image.fromarray(np.uint8(a2)).convert('L')
    #     a3 = Image.fromarray(np.uint8(a3)).convert('L')
    #     prediction = Image.merge('RGB', [a1, a2, a3]))  # 函数合并三个通道，此时prediction就成了三通道彩色图。
    #     prediction.save(test_save_path + '/' + case + '.png')

    # return metric_list

def testISIC_single_volume(image, label, net, classes, patch_size=[224, 224], test_save_path=None, z_spacing=1):
    image, label = image.squeeze(0).cpu().detach().numpy(), label.squeeze(0).cpu().detach().numpy()
    label = one_hot_encode(label, 3)
    if len(image.shape) == 3:
        prediction = np.zeros_like(label)  # 全零数组用来存储预测结果
        for ind in range(image.shape[0]):
            slice = image[ind, :, :]  # 选择当前切片
            x, y = slice.shape[0], slice.shape[1]  # 获取当前切片的高度和宽度
            if x != patch_size[0] or y != patch_size[1]:  # 如果当前切片的尺寸与patch_size不匹配，则调整切片的尺寸
                slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=3)  # previous using 0
                #  转换为张量，并对其添加额外的维度，并将其发送到GPU上。
            input = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0).float().cuda()
            net.eval()
            with torch.no_grad():
                outputs = net(input)
                out = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0)
                out = out.cpu().detach().numpy()
                if x != patch_size[0] or y != patch_size[1]:
                    pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0)
                else:
                    pred = out
                prediction[ind] = pred
    else:
        input = torch.from_numpy(image).unsqueeze(
            0).unsqueeze(0).float().cuda()
        net.eval()
        with torch.no_grad():
            out = torch.argmax(torch.softmax(net(input), dim=1), dim=1).squeeze(0)
            prediction = out.cpu().detach().numpy()
    metric_list = []
    for i in range(1, classes):
        metric_list.append(calculate_metric_percase(prediction == i, label == i))

    if test_save_path is not None:
        img_itk = sitk.GetImageFromArray(image.astype(np.float32))
        prd_itk = sitk.GetImageFromArray(prediction.astype(np.float32))
        lab_itk = sitk.GetImageFromArray(label.astype(np.float32))
        img_itk.SetSpacing((1, 1, z_spacing))
        prd_itk.SetSpacing((1, 1, z_spacing))
        lab_itk.SetSpacing((1, 1, z_spacing))
        sitk.WriteImage(prd_itk, test_save_path + '/' + "_pred.nii.gz")
        sitk.WriteImage(img_itk, test_save_path + '/' + "_img.nii.gz")
        sitk.WriteImage(lab_itk, test_save_path + '/' + "_gt.nii.gz")
    return metric_list

def one_hot_encode(label, num_channels):
    """
    生成标签的one-hot编码

    参数：
    label (ndarray): 输入的标签
    num_channels (int): 生成的通道数

    返回值：
    one_hot (ndarray): one-hot编码后的标签，形状为 (num_channels, label.shape[0], label.shape[1], ...)
    """
    label_shape = label.shape
    one_hot = np.zeros((num_channels,) + label_shape, dtype=np.float32)
    for channel in range(num_channels):
        one_hot[channel] = (label == channel).astype(np.float32)
    return one_hot
"""
    处理图片使用
"""
class myToTensor:
    def __init__(self):
        pass

    def __call__(self, data):
        image, mask = data
        return torch.tensor(image).permute(2, 0, 1), torch.tensor(mask).permute(2, 0, 1)

# 调整大小
class myResize:
    def __init__(self, size_h=256, size_w=256):
        self.size_h = size_h
        self.size_w = size_w

    def __call__(self, data):
        image, mask = data
        return TF.resize(image, [self.size_h, self.size_w]), TF.resize(mask, [self.size_h, self.size_w])

class myResize_npz:
    def __init__(self, x, y):
        self.size_h = x
        self.size_w = y

    def __call__(self, data):
        image, mask = data['image'], data['label']
        image, mask = torch.from_numpy(image), torch.from_numpy(mask)

        # image = np.array(image, dtype=np.float32)  # 将数据类型转换为 float32
        # image = (image * 255).astype(np.uint8)  # 将数据缩放到 0-255 范围，并转换为 uint8
        # # 将 Numpy 数组转换为 PIL 图像对象
        # image = Image.fromarray(image)
        # mask = Image.fromarray(mask)
        #
        # # 调整图像和掩码尺寸
        # image = image.resize((self.size_w, self.size_h))
        # mask = mask.resize((self.size_w, self.size_h))
        #
        # # 将 PIL 图像对象转换回 Numpy 数组
        # image = np.array(image)
        # mask = np.array(mask)

        # return image, mask
        return TF.resize(image, [self.size_h, self.size_w]), TF.resize(mask, [self.size_h, self.size_w])

# 随机水平翻转
class myRandomHorizontalFlip:
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, data):
        image, mask = data
        if random.random() < self.p:
            return TF.hflip(image), TF.hflip(mask)
        else:
            return image, mask

# 随机垂直翻转
class myRandomVerticalFlip:
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, data):
        image, mask = data
        if random.random() < self.p:
            return TF.vflip(image), TF.vflip(mask)
        else:
            return image, mask

# 随机旋转
class myRandomRotation:
    def __init__(self, p=0.5, degree=[0, 360]):
        self.angle = random.uniform(degree[0], degree[1])
        self.p = p

    def __call__(self, data):
        image, mask = data
        if random.random() < self.p:
            return TF.rotate(image, self.angle), TF.rotate(mask, self.angle)
        else:
            return image, mask

# 归一化
class myNormalize:
    def __init__(self, data_name, train=True):
        if data_name == 'isic':
            if train:
                self.mean = 157.561
                self.std = 26.706
            else:
                self.mean = 149.034
                self.std = 32.022
        elif data_name == 'ISIC2017':
            if train:
                self.mean = 159.922
                self.std = 28.871
            else:
                self.mean = 148.429
                self.std = 25.748
        elif data_name == 'cvcc':  # 还没调参数
            if train:
                self.mean =72.562195
                self.std = 54.454075
            else:
                self.mean = 72.562195
                self.std = 54.454075
        elif data_name == 'DSB2018':  # 还没调参数
            if train:
                self.mean = 159.922
                self.std = 28.871
            else:
                self.mean = 148.429
                self.std = 25.748
        else:
            if train:
                self.mean = 159.922
                self.std = 28.871
            else:
                self.mean = 148.429
                self.std = 25.748

    def __call__(self, data):
        img, msk = data
        img_normalized = (img - self.mean) / self.std
        img_normalized = ((img_normalized - np.min(img_normalized))
                          / (np.max(img_normalized) - np.min(img_normalized))) * 255.
        return img_normalized, msk


"""
    计算各种指标
confusionMetric  # 注意：此处横着代表预测值，竖着代表真实值，与之前介绍的相反
    P\L     P    N
    P      TP    FP
    N      FN    TN
"""


class SegmentationMetric(object):
    def __init__(self, numClass):
        self.numClass = numClass
        self.confusionMatrix = np.zeros((self.numClass,) * 2)

    def pixelAccuracy(self):
        # return all class overall pixel accuracy
        #  PA = acc = (TP + TN) / (TP + TN + FP + TN)
        acc = np.diag(self.confusionMatrix).sum() / self.confusionMatrix.sum()
        return acc

    def classPixelAccuracy(self):
        # return each category pixel accuracy(A more accurate way to call it precision)
        # acc = (TP) / TP + FP
        classAcc = np.diag(self.confusionMatrix) / self.confusionMatrix.sum(axis=1)
        return classAcc  # 返回的是一个列表值，如：[0.90, 0.80, 0.96]，表示类别1 2 3各类别的预测准确率

    def meanPixelAccuracy(self):
        classAcc = self.classPixelAccuracy()
        meanAcc = np.nanmean(classAcc)  # np.nanmean 求平均值，nan表示遇到Nan类型，其值取为0
        return meanAcc  # 返回单个值，如：np.nanmean([0.90, 0.80, 0.96, nan, nan]) = (0.90 + 0.80 + 0.96） / 3 =  0.89

    def meanIntersectionOverUnion(self):
        # Intersection = TP Union = TP + FP + FN
        # IoU = TP / (TP + FP + FN)
        intersection = np.diag(self.confusionMatrix)  # 取对角元素的值，返回列表
        union = np.sum(self.confusionMatrix, axis=1) + np.sum(self.confusionMatrix, axis=0) - np.diag(
            self.confusionMatrix)  # axis = 1表示混淆矩阵行的值，返回列表； axis = 0表示取混淆矩阵列的值，返回列表
        IoU = intersection / union  # 返回列表，其值为各个类别的IoU
        mIoU = np.nanmean(IoU)  # 求各类别IoU的平均
        return mIoU

    def genConfusionMatrix(self, imgPredict, imgLabel):  # 同FCN中score.py的fast_hist()函数
        # remove classes from unlabeled pixels in gt image and predict
        imgPredict,imgLabel = imgPredict.astype(int),imgLabel.astype(int)
        mask = (imgLabel >= 0) & (imgLabel < self.numClass)  # bool类型
        label = self.numClass * imgLabel[mask] + imgPredict[mask]  # ndarray float32

        # 后加
        # label_int = [int(x) for x in label]   # 或者根据判断改数值？
        label_int = [max(0, min(x, self.numClass - 1)) for x in label]  # bincount不能接收负数

        # 使用列表推导式将列表转换为整数类型
        count = np.bincount(label_int, minlength=self.numClass ** 2)  # shape为7，无法reshape
        confusionMatrix = count.reshape(self.numClass, self.numClass)
        return confusionMatrix

    def Frequency_Weighted_Intersection_over_Union(self):
        # FWIOU =     [(TP+FN)/(TP+FP+TN+FN)] *[TP / (TP + FP + FN)]
        freq = np.sum(self.confusion_matrix, axis=1) / np.sum(self.confusion_matrix)
        iu = np.diag(self.confusion_matrix) / (
                np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                np.diag(self.confusion_matrix))
        FWIoU = (freq[freq > 0] * iu[freq > 0]).sum()
        return FWIoU

    def addBatch(self, imgPredict, imgLabel):
        assert imgPredict.shape == imgLabel.shape
        self.confusionMatrix += self.genConfusionMatrix(imgPredict, imgLabel)

    def reset(self):
        self.confusionMatrix = np.zeros((self.numClass, self.numClass))


def mask2one_hot(label, num_classes):
   """
   label: 标签图像 # （batch_size, 1, h, w)
   out: 网络的输出
   """
   # num_classes = out.shape[1] # 分类类别数
   label = label.long()  # 将标签转换为整数类型

   current_label = label.squeeze(1)  #（batch_size, 1, h, w) ---> （batch_size, h, w)

   batch_size, h,w = current_label.shape[0],current_label.shape[1],current_label.shape[2]


   # print(h,w,batch_size)  # 224 224 24

   one_hots = []
   for i in range(num_classes):
       tmplate = torch.ones(batch_size, h, w) # （batch_size, h, w)
       tmplate[current_label != i] = 0
       tmplate = tmplate.view(batch_size,1,h,w) # （batch_size, h, w) --> （batch_size, 1, h, w)

       one_hots.append(tmplate)

   onehot = torch.cat(one_hots, dim=1)

   return onehot



def save_imgs(img, msk, msk_pred, i, save_path, datasets, threshold=0.5, test_data_name=None):
    img = img.squeeze(0).permute(1,2,0).detach().cpu().numpy()
    msk = msk.squeeze(0).permute(1,2,0).detach().cpu().numpy()
    # msk_pred = msk_pred.squeeze(0).permute(1,2,0).detach().cpu().numpy()
    img = img / 223. if img.max() > 1.1 else img
    if datasets == 'retinal':
        msk = np.squeeze(msk, axis=0)
        msk_pred = np.squeeze(msk_pred, axis=0)
    else:
        msk = np.where(np.squeeze(msk, axis=0) > 0.5, 1, 0)
        msk_pred = np.where(np.squeeze(msk_pred, axis=0) > threshold, 1, 0)

    plt.figure(figsize=(7,15))

    plt.subplot(3,1,1)
    plt.imshow(img)
    plt.axis('off')

    plt.subplot(3,1,2)
    plt.imshow(msk, cmap= 'gray')
    plt.axis('off')

    plt.subplot(3,1,3)
    plt.imshow(msk_pred, cmap = 'gray')
    plt.axis('off')

    if test_data_name is not None:
        save_path = save_path + test_data_name + '_'
    plt.savefig(save_path + str(i) +'.png')
    plt.close()




